import argparse
import os
from pytorch_grad_cam.guided_backprop import BaselineModel
from torchvision.transforms import Compose, Normalize, ToTensor
import cv2
import numpy as np
import torch
from torchvision import models

from pytorch_grad_cam import GradCAM, \
                             ScoreCAM, \
                             GradCAMPlusPlus, \
                             AblationCAM, \
                             XGradCAM, \
                             EigenCAM, \
                             EigenGradCAM

from pytorch_grad_cam import GuidedBackpropReLUModel
from pytorch_grad_cam.utils.image import show_cam_on_image, \
                                         deprocess_image, \
                                         preprocess_image


def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--use-cuda', type=str, default='0',
                        help='Use NVIDIA GPU acceleration, gpu list')
    parser.add_argument('--image-path', type=str, default='./examples/both.png',
                        help='Input image path')
    parser.add_argument('--aug_smooth', action='store_true',
                        help='Apply test time augmentation to smooth the CAM')
    parser.add_argument('--eigen_smooth', action='store_true',
                        help='Reduce noise by taking the first principle componenet'
                        'of cam_weights*activations')
    parser.add_argument('--method', type=str, default='gradcam',
                        choices=['gradcam', 'gradcam++', 'scorecam', 'xgradcam',
                                 'ablationcam', 'eigencam', 'eigengradcam'],
                        help='Can be gradcam/gradcam++/scorecam/xgradcam'
                             '/ablationcam/eigencam/eigengradcam')

    args = parser.parse_args()

    return args


if __name__ == '__main__':
    """ python cam.py -image-path <path_to_image>
    Example usage of loading an image, and computing:
        1. CAM
        2. Guided Back Propagation
        3. Combining both
    """

    args = get_args()
    use_cuda = args.use_cuda != '-1' and torch.cuda.is_available()
    if use_cuda:
        os.environ["CUDA_VISIBLE_DEVICES"] = args.use_cuda
        print('Using GPU for acceleration')
    else:
        print('Using CPU for computation')

    methods = \
        {"gradcam": GradCAM,
         "scorecam": ScoreCAM,
         "gradcam++": GradCAMPlusPlus,
         "ablationcam": AblationCAM,
         "xgradcam": XGradCAM,
         "eigencam": EigenCAM,
         "eigengradcam": EigenGradCAM}

    model = models.resnet50(pretrained=True)

    rgb_img = cv2.imread(args.image_path, 1)[:, :, ::-1]
    rgb_img = np.float32(rgb_img) / 255
    # input_tensor = preprocess_image(rgb_img, mean=[0.485, 0.456, 0.406], 
    #                                          std=[0.229, 0.224, 0.225])
    input_tensor = ToTensor()(rgb_img).unsqueeze(0)


    # If None, returns the map for the highest scoring category.
    # Otherwise, targets the requested category.
    target_category = None

    gb_model = GuidedBackpropReLUModel(model=model, use_cuda=use_cuda)
    # gb_model = BaselineModel(model=model, use_cuda=use_cuda)
    adv_img = gb_model.attack(input_img=input_tensor)

    ori_img = rgb_img * 255
    # cv2.imwrite(f'{args.method}_ori.jpg', ori_img[:, :, ::-1])

    cv2.imwrite(f'{args.method}_gb_attack.jpg', adv_img[:, :, ::-1])
    # cv2.imwrite(f'{args.method}_baseline_attack.jpg', adv_img[:, :, ::-1])

